import os
import cv2
import numpy as np
import tempfile
from scenedetect import VideoManager, SceneManager, StatsManager
from scenedetect.detectors import ContentDetector
from scenedetect.scene_manager import save_images


class VideoSceneSplitter:
    """视频场景切割工具类"""

    def __init__(self, threshold=30.0, min_scene_length=15):
        """
        初始化场景切割器

        参数:
            threshold (float): 场景变化检测阈值，较低的值会产生更多的场景分割
            min_scene_length (int): 最小场景长度(帧数)，防止过度分割
        """
        self.threshold = threshold
        self.min_scene_length = min_scene_length

    def detect_scenes(self, video_path):
        """
        检测视频中的场景变化点

        参数:
            video_path (str): 视频文件路径

        返回:
            list: 场景列表，每个场景为(开始帧,结束帧)元组
        """
        # 创建视频管理器
        video_manager = VideoManager([video_path])
        stats_manager = StatsManager()
        scene_manager = SceneManager(stats_manager)

        # 添加内容检测器
        scene_manager.add_detector(
            ContentDetector(threshold=self.threshold, min_scene_len=self.min_scene_length)
        )

        # 开始处理视频
        video_manager.start()

        # 检测场景
        scene_manager.detect_scenes(frame_source=video_manager)

        # 获取场景列表
        scene_list = scene_manager.get_scene_list()

        # 释放资源
        video_manager.release()

        return scene_list

    def extract_scenes(self, video_path, output_dir=None, format='mp4'):
        """
        提取视频中的所有场景到单独的视频文件

        参数:
            video_path (str): 视频文件路径
            output_dir (str): 输出目录，如不指定则使用视频所在目录
            format (str): 输出视频格式

        返回:
            list: 提取的场景视频文件路径列表
        """
        # 检查输出目录
        if output_dir is None:
            output_dir = os.path.dirname(video_path)

        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 获取视频名称(不含扩展名)
        video_name = os.path.splitext(os.path.basename(video_path))[0]

        # 检测场景
        scene_list = self.detect_scenes(video_path)

        # 获取视频帧率
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        cap.release()

        # 提取每个场景
        output_files = []
        for i, scene in enumerate(scene_list):
            start_frame = scene[0].get_frames()
            end_frame = scene[1].get_frames()

            # 转换帧号为时间(秒)
            start_time = start_frame / fps
            end_time = end_frame / fps
            duration = end_time - start_time

            # 输出文件名
            output_file = os.path.join(output_dir, f"{video_name}_scene_{i + 1:03d}.{format}")

            # 使用FFmpeg提取场景片段
            os.system(f'ffmpeg -i "{video_path}" -ss {start_time:.3f} -t {duration:.3f} '
                      f'-c:v libx264 -c:a aac -strict experimental "{output_file}" -y -loglevel warning')

            output_files.append(output_file)

        return output_files

    def extract_keyframes(self, video_path, output_dir=None, num_frames=1):
        """
        为每个场景提取关键帧图像

        参数:
            video_path (str): 视频文件路径
            output_dir (str): 输出目录，如不指定则使用视频所在目录
            num_frames (int): 每个场景提取的关键帧数量

        返回:
            dict: 场景索引到关键帧图像路径的映射
        """
        # 检查输出目录
        if output_dir is None:
            output_dir = os.path.dirname(video_path)

        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 获取视频名称(不含扩展名)
        video_name = os.path.splitext(os.path.basename(video_path))[0]

        # 检测场景
        video_manager = VideoManager([video_path])
        stats_manager = StatsManager()
        scene_manager = SceneManager(stats_manager)

        # 添加内容检测器
        scene_manager.add_detector(
            ContentDetector(threshold=self.threshold, min_scene_len=self.min_scene_length)
        )

        # 开始处理视频
        video_manager.start()

        # 检测场景
        scene_manager.detect_scenes(frame_source=video_manager)

        # 保存关键帧
        keyframe_paths = {}
        scene_list = scene_manager.get_scene_list()

        for i, scene in enumerate(scene_list):
            scene_keyframes_dir = os.path.join(output_dir, f"{video_name}_scene_{i + 1:03d}_keyframes")
            if not os.path.exists(scene_keyframes_dir):
                os.makedirs(scene_keyframes_dir)

            # 提取并保存关键帧
            keyframes = save_images(
                scene=scene,
                video_manager=video_manager,
                num_images=num_frames,
                output_dir=scene_keyframes_dir,
                image_name_template='$VIDEO_NAME-$SCENE_NUMBER-$IMAGE_NUMBER',
                suppress_output=True
            )

            keyframe_paths[i] = keyframes

        # 释放资源
        video_manager.release()

        return keyframe_paths

    def analyze_video_metadata(self, video_path):
        """
        分析视频元数据

        参数:
            video_path (str): 视频文件路径

        返回:
            dict: 视频元数据信息
        """
        cap = cv2.VideoCapture(video_path)

        metadata = {
            'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
            'fps': cap.get(cv2.CAP_PROP_FPS),
            'total_frames': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
            'duration': int(cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS))
        }

        cap.release()
        return metadata